import dolphindb as ddb
import pandas as pd
from datetime import datetime, timedelta

class ContractProvider:
    """合约管理器"""
    
    def __init__(self):
        self.session = None
        
    def connect_db(self):
        """连接DolphinDB数据库"""
        if self.session is None:
            self.session = ddb.session()
            self.session.connect("localhost", 8848, "admin", "123456")
            
    def close_db(self):
        """关闭数据库连接"""
        if self.session:
            self.session.close()
            self.session = None
            
    def get_main_contract_data(self, product_code, start_date=None, end_date=None):
        """获取主力合约数据
        
        Parameters
        ----------
        product_code : str
            品种代码，如 'rb'
        start_date : str, optional
            起始日期，格式：'YYYY-MM-DD'
        end_date : str, optional
            结束日期，格式：'YYYY-MM-DD'
        """
        try:
            self.connect_db()
            
            # 构建查询条件
            conditions = f"code like `{product_code}%"
            if start_date:
                conditions += f" and date >= date('{start_date}')"
            if end_date:
                conditions += f" and date <= date('{end_date}')"
            
            # 查询数据并按成交量排序找出每天的主力合约
            script = f"""
            t = select date, code, close, volume, amount
                from loadTable("dfs://futures", "daily_all")
                where {conditions};
                
            result = select date, code, close, volume, amount
                from t
                context by date
                having volume = max(volume)
                order by date;
            """
            
            return self.session.run(script)
            
        finally:
            self.close_db()
            
    def get_trading_contract_data(self, product_code, main_contract_code, months=[1,5,9], 
                                ahead_months=3, start_date=None, end_date=None):
        """获取实际交易合约数据
        
        Parameters
        ----------
        product_code : str
            品种代码，如 'rb'
        main_contract_code : str
            当前主力合约代码，如 'rb2401'
        months : list
            合约月份列表，默认[1,5,9]
        ahead_months : int
            提前几个月换仓，默认3个月
        start_date : str
            起始日期
        end_date : str
            结束日期
            
        Returns
        -------
        dict
            包含以下信息的字典：
            - current_contract: 当前应该交易的合约
            - next_contract: 下一个应该交易的合约
            - switch_date: 建议切换日期
            - contract_data: 相关合约的行情数据
        """
        try:
            self.connect_db()
            
            # 从主力合约代码中提取年月
            year = int(main_contract_code[-4:-2])
            month = int(main_contract_code[-2:])
            
            # 找出下一个交易的合约月份
            next_month = None
            for m in sorted(months):
                if m > month:
                    next_month = m
                    break
            if next_month is None:  # 如果当前是最后一个月份，则转到下一年第一个月
                next_month = months[0]
                year += 1
            
            # 构建实际交易合约代码
            current_contract = f"{product_code}{year:02d}{month:02d}"
            next_contract = f"{product_code}{year:02d}{next_month:02d}"
            
            # 计算建议切换日期
            switch_date = datetime(2000 + year, next_month, 1) - timedelta(days=ahead_months*30)
            
            # 获取这些合约的数据
            contracts = [current_contract, next_contract]
            conditions = f"code in {contracts}"
            if start_date:
                conditions += f" and date >= date('{start_date}')"
            if end_date:
                conditions += f" and date <= date('{end_date}')"
                
            script = f"""
            select date, code, open, high, low, close, volume, amount
            from loadTable("dfs://futures", "daily_all")
            where {conditions}
            order by date
            """
            
            contract_data = self.session.run(script)
            
            return {
                'current_contract': current_contract,
                'next_contract': next_contract,
                'switch_date': switch_date,
                'contract_data': contract_data
            }
            
        finally:
            self.close_db() 